#include <iostream>
#include <stdlib.h>

#include "gtest/gtest.h"

#include "kpextVkContext.h"
#include "KpExtGP.h"
#include "kpextOpImgCpyF32C1.h"

#include "kpextCtxGlobal.h"

using namespace std;

TEST(kpextFuncTest, OpImgCpyF32C1){
    ASSERT_NE(nullptr, GvkContextPtr);
    auto op = std::make_shared<kpext::OpImgCpyF32C1>(
                GvkContextPtr
                );

    auto cmdbuf = GvkContextPtr->cmdbuf();

    auto stagebuf1 = GvkContextPtr->stagingBuf(4 * sizeof (float));
    auto stagebuf2 = GvkContextPtr->stagingBuf(4 * sizeof (float));
    auto inbuf = GvkContextPtr->deviceBuf(4 * sizeof (float));
    auto outbuf = GvkContextPtr->deviceBuf(4 * sizeof (float));

    ASSERT_NE(nullptr, op);
    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, outbuf);

    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    {
        auto data = stagebuf1->data<float>();
        for (int i = 0; i < 4; ++i) {
            data[i] = static_cast<float>(i);
        }
        stagebuf1->unmap();

        data = stagebuf2->data<float>();
        for (int i = 0; i < 4; ++i) {
            data[i] = 99;
        }
        stagebuf2->unmap();
    }
    op->rebuild(4, 1,
                4, 1,
                inbuf,
                outbuf
                );

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    bufset.prepare(outbuf, 0x4479C000); // 这个也会视为是transfer stage
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    cmdbuf->record(&barrierRx);
    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&bufcpy);
    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();

    {
        auto data1 = stagebuf1->data<float>();
        auto data2 = stagebuf2->data<float>();

        for (int i = 0; i < 4; ++i) {
            EXPECT_EQ(data2[i], static_cast<float>(data1[i]));
        }

        stagebuf1->unmap();
        stagebuf2->unmap();
    }

    //       ------------------------

    int cols_1 = 50, rows_1 = 100;
    int cols_2 = rows_1, rows_2 = cols_1;
    auto size = static_cast<size_t>(cols_1 * rows_1);
    vector<float> initdata(size);
    for (size_t i = 0; i < initdata.size(); ++i) {
        initdata[i] = static_cast<float>(i);
    }

    size = size * sizeof (float);
    stagebuf1 = GvkContextPtr->stagingBuf(size);
    inbuf = GvkContextPtr->deviceBuf(size);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, inbuf);

    {
        auto dst = stagebuf1->data<float>();
        std::memcpy(dst, initdata.data(), size);
        stagebuf1->unmap();
    }


    size = static_cast<size_t>(cols_2 * rows_2) * sizeof (float);
    stagebuf2 = GvkContextPtr->stagingBuf(size);
    outbuf = GvkContextPtr->deviceBuf(size);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, outbuf);

    barrierTx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );


    op->rebuild(cols_1, rows_1,
                cols_2, rows_2,
                inbuf, outbuf
                );

    cmdbuf->begin();

    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    bufset.prepare(outbuf, 0x4479C000); // 这个也会视为是transfer stage
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    cmdbuf->record(&barrierRx);
    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&bufcpy);

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();

    int cols = (cols_2 < cols_1 ? cols_2 : cols_1);
    int rows = (rows_2 < rows_1 ? rows_2 : rows_1);
    bool isalleq = true;

    {
        float * p1 = stagebuf1->data<float>();
        float * p2 = stagebuf2->data<float>();
        for (int r = 0; r < rows; ++r) {
            float * sptr = p2;
            for (int c = 0; c < cols; ++c) {
                int idx1 = c + r * cols_1;
                int idx2 = c + r * cols_2;
                if(c == 0){
                    sptr = p2 + idx2;
                }
                float f1 = p1[idx1];
                float f2 = p2[idx2];
                isalleq = isalleq && ( f1 == f2);
            }
//            std::vector<float> rowdisp(sptr, sptr + cols);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf1->unmap();
        stagebuf2->unmap();

        //        SPDLOG_INFO(" {} ",stagebuf2->vector<float>());
    }

    EXPECT_EQ(true, isalleq);

    // ------------

    op->rebuild(cols_2, rows_2,
                cols_1, rows_1,
                outbuf, inbuf
                );

    cmdbuf->begin();

    bufcpy.prepare(outbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    bufset.prepare(inbuf, 0x4479C000); // 这个也会视为是transfer stage
    {
        std::memset(stagebuf2->data<void>(), 0, stagebuf2->getSize());
        stagebuf2->unmap();
    }
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    cmdbuf->record(&barrierRx);
    bufcpy.prepare(stagebuf2, inbuf);
    cmdbuf->record(&bufcpy);

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();

    cols = (cols_2 < cols_1 ? cols_2 : cols_1);
    rows = (rows_2 < rows_1 ? rows_2 : rows_1);
    isalleq = true;

    {
        float * p1 = stagebuf2->data<float>();
        float * p2 = stagebuf1->data<float>();
        for (int r = 0; r < rows; ++r) {
            float * sptr = p2;
            for (int c = 0; c < cols; ++c) {
                int idx1 = c + r * cols_1;
                int idx2 = c + r * cols_2;
                if(c == 0){
                    sptr = p2 + idx2;
                }
                float f1 = p1[idx1];
                float f2 = p2[idx2];
                isalleq = isalleq && ( f1 == f2);
            }
            //            std::vector<float> rowdisp(sptr, sptr + cols);
            //            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf1->unmap();
        stagebuf2->unmap();

        //        SPDLOG_INFO(" {} ",stagebuf2->vector<float>());
    }

    EXPECT_EQ(true, isalleq);

}

GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");

    testing::InitGoogleTest(&argc, argv);

    testing::Environment* const foo_env =
            testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    UNUSE_VAR(foo_env);

    return RUN_ALL_TESTS();
}
